from ase.build import bulk
from torch_sim.models.orb import OrbModel
from torch_sim.state import initialize_state
from orb_models.forcefield import pretrained
import torch
import os
import torch_sim as ts
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
orbff = pretrained.orb_v3_conservative_inf_omat(
    device=device,
    precision="float32-highest",
    compile=False
)
orb_model = OrbModel(model=orbff, dtype=torch.get_default_dtype())

from torch_sim.integrators.nvt import nvt_langevin

n_steps = 30
timestep = 0.001
temperature = 300.0  # K
filename = "nvt_trajectory.h5"
save_dir = "notebook/results2"

filename = os.path.join(save_dir, filename)
reporter = ts.TrajectoryReporter(
    filenames=[filename],
    state_frequency=10,
)

system = bulk('Ti')
system = initialize_state(system, device=device, dtype=torch.get_default_dtype())
print(orb_model(system)["energy"]) # tensor(-15.5739, device='cuda:0', grad_fn=<SqueezeBackward0>)

# Run NVT simulation
_ = ts.integrate(
    system,
    orb_model,
    # sevennet_model,
    integrator=nvt_langevin,
    n_steps=n_steps,
    timestep=timestep,
    temperature=temperature,
    trajectory_reporter=reporter,
)
## RuntimeError: split expects at least a 1-dimensional tensor